import cv2.aruco as aruco
from PIL import Image
import datetime
import _datetime
import itertools
import numpy as np
import cv2
import time
import math
import threading
import sys
from fps_limiter import LimitFPS
from dji_asdk_to_python.errors import CustomError
from dji_asdk_to_python.flight_controller.virtual_stick.flight_control_data import FlightControlData
from dji_asdk_to_python.flight_controller.flight_controller_state import FlightControllerState
from dji_asdk_to_python.camera.exposure_mode import ExposureMode
from dji_asdk_to_python.camera.iso import ISO
from dji_asdk_to_python.camera.shutter_speed import ShutterSpeed
from dji_asdk_to_python.utils.FPS import FPS
from dji_asdk_to_python.flight_controller.virtual_stick.control_mode import (
    VerticalControlMode
)
from dji_asdk_to_python.products.aircraft import Aircraft
from dji_asdk_to_python.precision_landing.landing import ArucoLanding


"""
This demo calculates multiple things for different scenarios.
IF RUNNING ON A PI, BE SURE TO sudo modprobe bcm2835-v4l2
Here are the defined reference frames:
TAG:
                A y
                |
                |
                |tag center
                O---------> x
CAMERA:
                X--------> x
                | frame center
                |
                |
                V y
F1: Flipped (180 deg) tag frame around x axis
F2: Flipped (180 deg) camera frame around x axis
The attitude of a generic frame 2 respect to a frame 1 can obtained by calculating euler(R_21.T)
We are going to obtain the following quantities:
    > from aruco library we obtain tvec and Rct, position of the tag in camera frame and attitude of the tag
    > position of the Camera in Tag axis: -R_ct.T*tvec
    > Transformation of the camera, respect to f1 (the tag flipped frame): R_cf1 = R_ct*R_tf1 = R_cf*R_f
    > Transformation of the tag, respect to f2 (the camera flipped frame): R_tf2 = Rtc*R_cf2 = R_tc*R_f
    > R_tf1 = R_cf2 an symmetric = R_f
"""


class ArucoSingleTracker:
    def __init__(self, camera_matrix, camera_distortion):

        self._camera_matrix = camera_matrix
        self._camera_distortion = camera_distortion

        # --- 180 deg rotation matrix around the x axis
        self._R_flip = np.zeros((3, 3), dtype=np.float32)
        self._R_flip[0, 0] = 1.0
        self._R_flip[1, 1] = -1.0
        self._R_flip[2, 2] = -1.0

        # --- Define the aruco dictionary
        self._aruco_dict = aruco.custom_dictionary_from(
            20, 4, aruco.getPredefinedDictionary(aruco.DICT_4X4_100)
        )

        self._parameters = aruco.DetectorParameters_create()

    def _rotationMatrixToEulerAngles(self, R):
        # Calculates rotation matrix to euler angles
        # The result is the same as MATLAB except the order
        # of the euler angles ( x and z are swapped ).

        def isRotationMatrix(R):
            Rt = np.transpose(R)
            shouldBeIdentity = np.dot(Rt, R)
            I = np.identity(3, dtype=R.dtype)
            n = np.linalg.norm(I - shouldBeIdentity)
            return n < 1e-6

        assert isRotationMatrix(R)

        sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])

        singular = sy < 1e-6

        if not singular:
            x = math.atan2(R[2, 1], R[2, 2])
            y = math.atan2(-R[2, 0], sy)
            z = math.atan2(R[1, 0], R[0, 0])
        else:
            x = math.atan2(-R[1, 2], R[1, 1])
            y = math.atan2(-R[2, 0], sy)
            z = 0

        return np.array([x, y, z])

    def track(
        self, frame, id_to_find=None, marker_size=None,
    ):

        marker_found = False
        x = y = z = pitch_camera = x_camera = y_camera = z_camera = 0

        # -- Convert in gray scale
        gray = cv2.cvtColor(
            frame, cv2.COLOR_BGR2GRAY
        )  # -- remember, OpenCV stores color images in Blue, Green, Red

        # -- Find all the aruco markers in the image
        corners, ids, rejected = aruco.detectMarkers(
            image=gray,
            dictionary=self._aruco_dict,
            parameters=self._parameters,
            cameraMatrix=self._camera_matrix,
            distCoeff=self._camera_distortion,
        )
        pitch_marker, roll_marker, yaw_marker = None, None, None
        pitch_camera, roll_camera, yaw_camera = None, None, None

        planned_ids = []
        if ids is not None:
            planned_ids = list(itertools.chain(*ids))
        if id_to_find in planned_ids:
            index_id_to_find = planned_ids.index(id_to_find)
            marker_found = True
            # -- array of rotation and position of each marker in camera frame
            # -- rvec = [[rvec_1], [rvec_2], ...]    attitude of the marker respect to camera frame
            # -- tvec = [[tvec_1], [tvec_2], ...]    position of the marker in camera frame
            rvecs, tvecs, _ = aruco.estimatePoseSingleMarkers(
                corners, marker_size, self._camera_matrix, self._camera_distortion
            )

            # -- Unpack the output
            rvec, tvec = rvecs[index_id_to_find][0], tvecs[index_id_to_find][0]

            x = tvec[0]
            y = tvec[1]
            z = tvec[2]

            # -- Draw the detected marker and put a reference frame over it
            aruco.drawDetectedMarkers(frame, corners)
            aruco.drawAxis(frame, camera_matrix, camera_distortion, rvec, tvec, 10)

            # -- Print the tag position in camera frame
            str_position = "MARKER Position x=%4.0f  y=%4.0f  z=%4.0f" % (tvec[0], tvec[1], tvec[2])
            cv2.putText(frame, str_position, (0, 100), font, 1, (0, 255, 0), 2, cv2.LINE_AA)

            # -- Obtain the rotation matrix tag->camera
            R_ct = np.matrix(cv2.Rodrigues(rvec)[0])
            R_tc = R_ct.T

            # -- Obtain the rotation matrix tag->camera
            R_ct = np.matrix(cv2.Rodrigues(rvec)[0])
            R_tc = R_ct.T

            # -- Get the attitude in terms of euler 321 (Needs to be flipped first)
            (
                roll_marker,
                pitch_marker,
                yaw_marker,
            ) = self._rotationMatrixToEulerAngles(self._R_flip * R_tc)

            # -- Print the marker's attitude respect to camera frame
            str_attitude = "MARKER Attitude r=%4.0f  p=%4.0f  y=%4.0f" % (
            math.degrees(roll_marker), math.degrees(pitch_marker),
            math.degrees(yaw_marker))
            cv2.putText(frame, str_attitude, (0, 150), font, 1, (0, 255, 0), 2, cv2.LINE_AA)

            # -- Now get Position and attitude f the camera respect to the marker
            pos_camera = -R_tc * np.matrix(tvec).T
            x_camera = pos_camera[0]
            y_camera = pos_camera[1]
            z_camera = pos_camera[2]

            (
                roll_camera,
                pitch_camera,
                yaw_camera,
            ) = self._rotationMatrixToEulerAngles(self._R_flip * R_tc)

        if yaw_marker is None:
            marker_found = False
            yaw_marker = 0

        if marker_found:
            roll_camera = math.degrees(roll_camera)
            yaw_camera = math.degrees(yaw_camera)
            pitch_camera = math.degrees(pitch_camera)
            roll_marker = math.degrees(roll_marker)
            yaw_marker = math.degrees(yaw_marker)
            pitch_marker = math.degrees(pitch_marker)
            x_camera = float(x_camera)
            y_camera = float(y_camera)
            z_camera = float(z_camera)

        result = (
            marker_found,
            x,
            y,
            z,
            x_camera,
            y_camera,
            z_camera,
            roll_marker,
            yaw_marker,
            pitch_marker,
            roll_marker,
            roll_camera,
            yaw_camera,
            pitch_camera,
        )
        return result


class ArucoLanding:
    LANDING_CM = 275
    SECONDS_BEFORE_GET_UP = 10
    MAX_SECONDS_GETTING_LANDING = 15
    X_Y_CM_ERROR_ALLOWED = 12
    YAW_ERROR_ALLOWED = 15

    CURRENT_ISO = "PROGRAM"

    """
    Inits aruco precision landing
        Parameters:
        drone_ip (str) -> The IP of the drone
        camera_matrix (ndarray) -> The camera matrix of the drone's camera
        camera_distortion (ndarray) -> The camera distortion of the drone's camera
        marker_id (int) -> The ID of the aruco marker to be detected on the landing stage
        marker_size_cm (int) -> The size in CM of the aruco marker to be detected in the stage
    """

    def resetPid(self):
        self.p = 0.004
        self.i = 0.000005
        self.d = 0.0005

        self.pidx = PID(P=self.p, I=self.i, D=self.d)
        self.pidy = PID(P=self.p, I=self.i, D=self.d)
        self.pidz = PID(P=self.p, I=self.i, D=self.d)

        self.pidx.SetPoint = 0.0
        self.pidy.SetPoint = 15.0
        self.pidz.SetPoint = ArucoLanding.LANDING_CM / 1.2

        self.pidx.setSampleTime(0.1)
        self.pidy.setSampleTime(0.1)
        self.pidz.setSampleTime(0.1)

    def __init__(self, aircraft, marker_id, marker_size_cm, width, height, camera_matrix, camera_distortion):
        self.aircraft = aircraft
        self.cv2_manager = self.aircraft.getLiveStreamManager().getCV2Manager(with_buffer=True)
        self.cv2_manager.setWidth(width)
        self.cv2_manager.setHeigth(height)

        self.marker_id = marker_id
        self.marker_size_cm = marker_size_cm
        self.ast = ArucoSingleTracker(camera_distortion=camera_distortion, camera_matrix=camera_matrix)
        self.resetPid()

        self.running = True
        t2 = threading.Thread(target=self.revive_virtual_stick, args=[])
        t2.start()

    def camera_iso_setup(self, camera, cameraType):

        camera.setExposureMode(ExposureMode.MANUAL)

        if cameraType == "DAY_VERY_SUNNY":
            camera.setISO(ISO.ISO_50)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_8000)
        elif cameraType == "DAY_SUNNY":
            camera.setISO(ISO.ISO_200)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_4000)
        elif cameraType == "DAY_AFTERNOON":
            camera.setISO(ISO.ISO_400)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_2000)
        elif cameraType == "EVENING":
            camera.setISO(ISO.ISO_400)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_8000)
        elif cameraType == "MORNING":
            camera.setISO(ISO.ISO_400)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_4000)
        else:
            camera.setExposureMode(ExposureMode.PROGRAM)

    def camera_iso_print(self, camera, cameraType):

        camera.setExposureMode(ExposureMode.MANUAL)

        if cameraType == "DAY_VERY_SUNNY":
            camera.setISO(ISO.ISO_50)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_8000)
        elif cameraType == "DAY_SUNNY":
            camera.setISO(ISO.ISO_200)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_4000)
        elif cameraType == "DAY_AFTERNOON":
            camera.setISO(ISO.ISO_400)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_2000)
        elif cameraType == "EVENING":
            camera.setISO(ISO.ISO_800)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_2000)
        elif cameraType == "MORNING":
            camera.setISO(ISO.ISO_400)
            camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_4000)
        else:
            camera.setExposureMode(ExposureMode.PROGRAM)

    def start(self, is_night):
        result = self.cv2_manager.startStream("9003")
        if isinstance(result, CustomError):
            raise Exception("%s" % result)

        gimbal = self.aircraft.getGimbal()
        gimbal.rotate(-90, 0, 0)

        fc = self.aircraft.getFlightController()
        fc.setVirtualStickModeEnabled(True)
        fcd = FlightControlData()
        fc.setCollisionAvoidanceEnabled(False)
        fc.setVerticalControlMode(VerticalControlMode.VELOCITY)

        camera = self.aircraft.getCamera()

        print("LANDING IS NIGHT", is_night)

        camera.setExposureMode(ExposureMode.MANUAL)

        cameraTypeSettings = ["DAY_VERY_SUNNY", "DAY_SUNNY", "DAY_AFTERNOON", "EVENING", "MORNING", "PROGRAM"]
        # cameraTypeSettings = ["DAY_SUNNY","DAY_AFTERNOON"]
        # if is_night:
        #     camera.setExposureMode(ExposureMode.PROGRAM)
        #     camera.setISO(ISO.ISO_100)
        # else:
        #     camera.setExposureMode(ExposureMode.MANUAL)
        #     camera.setISO(ISO.ISO_800)
        #     camera.setShutterSpeed(ShutterSpeed.SHUTTER_SPEED_1_8000)

        start = time.perf_counter()
        last_z = sys.maxsize
        last_x = sys.maxsize
        last_y = sys.maxsize
        fps = FPS()
        fps_limiter = LimitFPS(fps=15)
        i = 0
        isoCountChanger = 0
        maxChance = 0
        rightIso = False
        self.camera_iso_setup(camera, cameraTypeSettings[i])

        # LOGGING

        # file1 = open("log"+str(datetime.date.today().day)+str(datetime.date.today().month)+str(datetime.date.today().year)+str(datetime.datetime.now().hour)+str(datetime.datetime.now().minute),"w")
        # file1.write("PRECISION,X,Y,Z,EXPOSURE_MODE,ISO,SHUTTER_SPEED"+"\n")
        # file1.write("-----------------------------------------------"+"\n")
        while True:
            end = time.perf_counter()

            frame = self.cv2_manager.getFrame()
            if maxChance <= 100:
                maxChance = maxChance + 1

            if frame is None:
                print("FRAME IS NONE")
                continue

            # --- Display the frame
            cv2.imshow('frame', frame)

            (
                marker_found,
                x_marker,
                y_marker,
                z_marker,
                x_camera,
                y_camera,
                z_camera,
                roll_marker,
                yaw_marker,
                pitch_marker,
                roll_marker,
                roll_camera,
                yaw_camera,
                pitch_camera,
            ) = self.ast.track(frame, self.marker_id, self.marker_size_cm)

            print("1 PRE MARKER FOUND")

            if marker_found and not rightIso:
                isoCountChanger = isoCountChanger + 1
                if isoCountChanger >= 50:
                    print("FOUND RIGHT ISO. PRECISION GREATER THAN 50%")
                    # file1.write(str(isoCountChanger)+","+str(last_x)+","+str(last_y)+","+str(last_z)+","+str(camera.getExposureMode())+","+str(camera.getISO())+","+str(camera.getShutterSpeed())+"\n")
                    rightIso = True
            elif marker_found and rightIso:
                print("2 MARKER FOUND")

                print("4 FPS marker detection %s" % (fps_))
                print("x %s y %s z %s yaw %s" % (x_marker, y_marker, z_marker, yaw_camera))

            else:
                print("2 MARKER NOT FOUND, LAST Z ", last_z)
                print("TIME BEFORE GOING UP ", (end - start))

                im = Image.fromarray(frame)
                im.save(
                    "/home/aras/aras-current/aras-control-service/Images/markeNotFound" + datetime.datetime.now().strftime(
                        '%Y_%m_%d_%H_%M_%S_%f') + ".jpeg")


                if not rightIso and maxChance > 100:
                    # if last_x == sys.maxsize or last_y == sys.maxsize or last_z == sys.maxsize:
                    #     file1.write(str(isoCountChanger)+",N/A,N/A,N/A,"+str(camera.getExposureMode())+","+str(camera.getISO())+","+str(camera.getShutterSpeed())+"\n")
                    # else:
                    #     file1.write(str(isoCountChanger)+","+str(last_x)+","+str(last_y)+","+str(last_z)+","+str(camera.getExposureMode())+","+str(camera.getISO())+","+str(camera.getShutterSpeed())+"\n")

                    print("#### CHANGING ISO #### --- PRECISION: ", str(isoCountChanger))
                    maxChance = 0
                    isoCountChanger = 0
                    i = i + 1
                    if i < len(cameraTypeSettings):
                        self.camera_iso_setup(camera, cameraTypeSettings[i])
                    else:
                        i = 0
                        self.camera_iso_setup(camera, cameraTypeSettings[i])


APP_IP = "192.168.20.37"
MARKER_ID = 17
MARKER_SIZE_CM = 60
IMAGE_WIDTH = 1280
IMAGE_HEIGHT = 720
CAMERA_MATRIX = np.loadtxt("calibration/camera_matrix.txt", delimiter=",")
CAMERA_DISTORTION = np.loadtxt("calibration/camera_distortion.txt", delimiter=",")

aircraft = Aircraft(APP_IP)
aruco_landing = ArucoLanding(aircraft=aircraft,
                             width=IMAGE_WIDTH,
                             height=IMAGE_HEIGHT,
                             camera_distortion=CAMERA_DISTORTION,
                             camera_matrix=CAMERA_MATRIX,
                             marker_id=MARKER_ID,
                             marker_size_cm=MARKER_SIZE_CM)
aruco_landing.start(is_night=False)
